{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e3e7ac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | default_exp _testing.in_memory_broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78dd1499",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "import asyncio\n",
    "import hashlib\n",
    "import random\n",
    "import string\n",
    "import uuid\n",
    "from contextlib import contextmanager\n",
    "from dataclasses import dataclass, field\n",
    "from typing import *\n",
    "from types import ModuleType\n",
    "\n",
    "from aiokafka.structs import ConsumerRecord, RecordMetadata, TopicPartition\n",
    "\n",
    "import fastkafka._components.aiokafka_consumer_loop\n",
    "import fastkafka._aiokafka_imports\n",
    "from fastkafka._aiokafka_imports import AIOKafkaConsumer, AIOKafkaProducer\n",
    "from fastkafka._components.logger import get_logger\n",
    "from fastkafka._components.meta import (\n",
    "    _get_default_kwargs_from_sig,\n",
    "    classcontextmanager,\n",
    "    delegates,\n",
    "    patch,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df56c33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import unittest\n",
    "from contextlib import asynccontextmanager\n",
    "\n",
    "import pytest\n",
    "\n",
    "from fastkafka.testing import ApacheKafkaBroker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e65cf05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "logger = get_logger(__name__)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74e8c6f7",
   "metadata": {},
   "source": [
    "# Local Kafka broker\n",
    "> In-memory mockup of Kafka broker protocol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96928f97",
   "metadata": {},
   "source": [
    "## Kafka partition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a64d6bc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class KafkaRecord:\n",
    "    topic: str = \"\"\n",
    "    partition: int = 0\n",
    "    key: Optional[bytes] = None\n",
    "    value: bytes = b\"\"\n",
    "    offset: int = 0\n",
    "    timestamp = 0\n",
    "    timestamp_type = 0\n",
    "    checksum = 0\n",
    "    serialized_key_size = 0\n",
    "    serialized_value_size = 0\n",
    "    headers: Sequence[Tuple[str, bytes]] = field(default_factory=list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1fae851",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class KafkaPartition:\n",
    "    def __init__(self, *, partition: int, topic: str):\n",
    "        \"\"\"\n",
    "        Initialize a KafkaPartition object.\n",
    "\n",
    "        Args:\n",
    "            partition: The partition number.\n",
    "            topic: The topic name.\n",
    "        \"\"\"\n",
    "        self.partition = partition\n",
    "        self.topic = topic\n",
    "        self.messages: List[KafkaRecord] = list()\n",
    "\n",
    "    def write(self, value: bytes, key: Optional[bytes] = None) -> RecordMetadata:  # type: ignore\n",
    "        \"\"\"\n",
    "        Write a Kafka record to the partition.\n",
    "\n",
    "        Args:\n",
    "            value: The value of the record.\n",
    "            key: The key of the record.\n",
    "\n",
    "        Returns:\n",
    "            The record metadata.\n",
    "        \"\"\"\n",
    "        record = KafkaRecord(\n",
    "            topic=self.topic,\n",
    "            partition=self.partition,\n",
    "            value=value,\n",
    "            key=key,\n",
    "            offset=len(self.messages),\n",
    "        )\n",
    "        record_meta = RecordMetadata(\n",
    "            topic=self.topic,\n",
    "            partition=self.partition,\n",
    "            topic_partition=TopicPartition(topic=self.topic, partition=self.partition),\n",
    "            offset=len(self.messages),\n",
    "            timestamp=1680602752070,\n",
    "            timestamp_type=0,\n",
    "            log_start_offset=0,\n",
    "        )\n",
    "        self.messages.append(record)\n",
    "        return record_meta\n",
    "\n",
    "    def read(self, offset: int) -> Tuple[List[KafkaRecord], int]:\n",
    "        \"\"\"\n",
    "        Read Kafka records from the partition starting from the given offset.\n",
    "\n",
    "        Args:\n",
    "            offset: The starting offset.\n",
    "\n",
    "        Returns:\n",
    "            A tuple containing the list of records and the current offset.\n",
    "        \"\"\"\n",
    "        return self.messages[offset:], len(self.messages)\n",
    "\n",
    "    def latest_offset(self) -> int:\n",
    "        \"\"\"\n",
    "        Get the latest offset of the partition.\n",
    "\n",
    "        Returns:\n",
    "            The latest offset.\n",
    "        \"\"\"\n",
    "        return len(self.messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99e5b162",
   "metadata": {},
   "outputs": [],
   "source": [
    "partition_index = 0\n",
    "topic = \"test\"\n",
    "partition = KafkaPartition(partition=partition_index, topic=topic)\n",
    "\n",
    "msgs = [b\"some_msg\" for _ in range(25)]\n",
    "\n",
    "expected = [\n",
    "    KafkaRecord(topic=topic, partition=partition_index, value=msg, offset=offset)\n",
    "    for offset, msg in enumerate(msgs)\n",
    "]\n",
    "\n",
    "for msg in msgs:\n",
    "    partition.write(msg)\n",
    "\n",
    "for offset in [0, 10, 20]:\n",
    "    actual = partition.read(offset=offset)\n",
    "\n",
    "    assert actual == (expected[offset:], len(msgs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be209005",
   "metadata": {},
   "outputs": [],
   "source": [
    "partition_index = 0\n",
    "topic = \"test\"\n",
    "key = b\"some_key\"\n",
    "partition = KafkaPartition(partition=partition_index, topic=topic)\n",
    "\n",
    "msgs = [b\"some_msg\" for _ in range(25)]\n",
    "expected = [\n",
    "    KafkaRecord(\n",
    "        topic=topic, partition=partition_index, value=msg, key=key, offset=offset\n",
    "    )\n",
    "    for offset, msg in enumerate(msgs)\n",
    "]\n",
    "\n",
    "for msg in msgs:\n",
    "    partition.write(msg, key=key)\n",
    "\n",
    "for offset in [0, 10, 20]:\n",
    "    actual = partition.read(offset=offset)\n",
    "\n",
    "    assert actual == (expected[offset:], len(msgs)), print(f\"{actual} != {expected}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "145cd897",
   "metadata": {},
   "source": [
    "## Kafka topic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b02c802e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class KafkaTopic:\n",
    "    def __init__(self, topic: str, num_partitions: int = 1):\n",
    "        \"\"\"\n",
    "        Initialize a KafkaTopic object.\n",
    "\n",
    "        Args:\n",
    "            topic: The topic name.\n",
    "            num_partitions: The number of partitions in the topic (default: 1).\n",
    "        \"\"\"\n",
    "        self.topic = topic\n",
    "        self.num_partitions = num_partitions\n",
    "        self.partitions: List[KafkaPartition] = [\n",
    "            KafkaPartition(topic=topic, partition=partition_index)\n",
    "            for partition_index in range(num_partitions)\n",
    "        ]\n",
    "\n",
    "    def read(  # type: ignore\n",
    "        self, partition: int, offset: int\n",
    "    ) -> Tuple[TopicPartition, List[KafkaRecord], int]:\n",
    "        \"\"\"\n",
    "        Read records from the specified partition and offset.\n",
    "\n",
    "        Args:\n",
    "            partition: The partition index.\n",
    "            offset: The offset from which to start reading.\n",
    "\n",
    "        Returns:\n",
    "            A tuple containing the topic partition, list of Kafka records, and the new offset.\n",
    "        \"\"\"\n",
    "        topic_partition = TopicPartition(topic=self.topic, partition=partition)\n",
    "        records, offset = self.partitions[partition].read(offset)\n",
    "        return topic_partition, records, offset\n",
    "\n",
    "    def write_with_partition(  # type: ignore\n",
    "        self,\n",
    "        value: bytes,\n",
    "        partition: int,\n",
    "    ) -> RecordMetadata:\n",
    "        \"\"\"\n",
    "        Write a record with a specified partition.\n",
    "\n",
    "        Args:\n",
    "            value: The value of the record.\n",
    "            partition: The partition to write the record to.\n",
    "\n",
    "        Returns:\n",
    "            The metadata of the written record.\n",
    "        \"\"\"\n",
    "        return self.partitions[partition].write(value)\n",
    "\n",
    "    def write_with_key(self, value: bytes, key: bytes) -> RecordMetadata:  # type: ignore\n",
    "        \"\"\"\n",
    "        Write a record with a specified key.\n",
    "\n",
    "        Args:\n",
    "            value: The value of the record.\n",
    "            key: The key of the record.\n",
    "\n",
    "        Returns:\n",
    "            The metadata of the written record.\n",
    "        \"\"\"\n",
    "        partition = int(hashlib.sha256(key).hexdigest(), 16) % self.num_partitions\n",
    "        return self.partitions[partition].write(value, key=key)\n",
    "\n",
    "    def write(  # type: ignore\n",
    "        self,\n",
    "        value: bytes,\n",
    "        *,\n",
    "        key: Optional[bytes] = None,\n",
    "        partition: Optional[int] = None,\n",
    "    ) -> RecordMetadata:\n",
    "        \"\"\"\n",
    "        Write a record to the topic.\n",
    "\n",
    "        Args:\n",
    "            value: The value of the record.\n",
    "            key: The key of the record (optional).\n",
    "            partition: The partition to write the record to (optional).\n",
    "\n",
    "        Returns:\n",
    "            The metadata of the written record.\n",
    "        \"\"\"\n",
    "        if partition is not None:\n",
    "            return self.write_with_partition(value, partition)\n",
    "\n",
    "        if key is not None:\n",
    "            return self.write_with_key(value, key)\n",
    "\n",
    "        partition = random.randint(0, self.num_partitions - 1)  # nosec\n",
    "        return self.write_with_partition(value, partition)\n",
    "\n",
    "    def latest_offset(self, partition: int) -> int:\n",
    "        \"\"\"\n",
    "        Get the latest offset of a partition.\n",
    "\n",
    "        Args:\n",
    "            partition: The partition index.\n",
    "\n",
    "        Returns:\n",
    "            The latest offset of the partition.\n",
    "        \"\"\"\n",
    "        return self.partitions[partition].latest_offset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1fcfe86",
   "metadata": {},
   "outputs": [],
   "source": [
    "msg = b\"msg\"\n",
    "\n",
    "topic = KafkaTopic(\"test_topic\", 1)\n",
    "\n",
    "expected = RecordMetadata(\n",
    "    topic=\"test_topic\",\n",
    "    partition=0,\n",
    "    topic_partition=TopicPartition(topic=\"test_topic\", partition=0),\n",
    "    offset=0,\n",
    "    timestamp=1680602752070,\n",
    "    timestamp_type=0,\n",
    "    log_start_offset=0,\n",
    ")\n",
    "actual = topic.write(msg)\n",
    "\n",
    "assert expected == actual\n",
    "\n",
    "expected = RecordMetadata(\n",
    "    topic=\"test_topic\",\n",
    "    partition=0,\n",
    "    topic_partition=TopicPartition(topic=\"test_topic\", partition=0),\n",
    "    offset=1,\n",
    "    timestamp=1680602752070,\n",
    "    timestamp_type=0,\n",
    "    log_start_offset=0,\n",
    ")\n",
    "actual = topic.write(msg, key=b\"123\")\n",
    "\n",
    "assert expected == actual, actual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb4ba74",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_name = \"test_topic\"\n",
    "msgs = [b\"msg\" for _ in range(1000)]\n",
    "partition_num = 10\n",
    "\n",
    "topic = KafkaTopic(topic_name, partition_num)\n",
    "\n",
    "# write to topic\n",
    "for msg in msgs:\n",
    "    topic.write(msg)\n",
    "\n",
    "# For each partition in topic check:\n",
    "for partition in range(partition_num):\n",
    "    topic_partition_expected = TopicPartition(topic=topic_name, partition=partition)\n",
    "    topic_partition_actual, data, _ = topic.read(partition=partition, offset=0)\n",
    "\n",
    "    # Read returns correct TopicPartition key\n",
    "    assert topic_partition_actual == topic_partition_expected\n",
    "\n",
    "    # Data is written into partition\n",
    "    assert len(data) > 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f492339",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_name = \"test_topic\"\n",
    "msgs = [b\"msg\" for _ in range(1000)]\n",
    "partition_num = 2\n",
    "\n",
    "topic = KafkaTopic(topic_name, partition_num)\n",
    "\n",
    "# write to topic with defined partition\n",
    "for msg in msgs:\n",
    "    topic.write(msg, partition=0)\n",
    "\n",
    "lengths = [len(topic.read(partition=i, offset=0)[1]) for i in range(partition_num)]\n",
    "\n",
    "assert [1000, 0] == lengths"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fff42583",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_name = \"test_topic\"\n",
    "msgs = [b\"msg\" for _ in range(1000)]\n",
    "partition_num = 3\n",
    "\n",
    "topic = KafkaTopic(topic_name, partition_num)\n",
    "\n",
    "# write to topic with defined key\n",
    "for msg in msgs[:450]:\n",
    "    topic.write(msg, key=b\"some_key\")\n",
    "\n",
    "for msg in msgs[450:]:\n",
    "    topic.write(msg, key=b\"some_key443\")\n",
    "\n",
    "lengths = [len(topic.read(partition=i, offset=0)[1]) for i in range(partition_num)]\n",
    "\n",
    "assert [0, 450, 550] == sorted(lengths)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0280992",
   "metadata": {},
   "source": [
    "## Group metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1f0e7c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "def split_list(list_to_split: List[Any], split_size: int) -> List[List[Any]]:\n",
    "    \"\"\"\n",
    "    Split a list into smaller lists of a specified size.\n",
    "\n",
    "    Args:\n",
    "        list_to_split: The list to split.\n",
    "        split_size: The size of each split.\n",
    "\n",
    "    Returns:\n",
    "        A list of smaller lists.\n",
    "    \"\"\"\n",
    "    return [\n",
    "        list_to_split[start_index : start_index + split_size]\n",
    "        for start_index in range(0, len(list_to_split), split_size)\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d58bb33",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert split_list([1, 2, 3, 4, 5], 1) == [[1], [2], [3], [4], [5]]\n",
    "assert split_list([1, 2, 3, 4, 5], 2) == [[1, 2], [3, 4], [5]]\n",
    "assert split_list([1, 2, 3, 4, 5], 3) == [[1, 2, 3], [4, 5]]\n",
    "assert split_list([1, 2, 3, 4, 5], 5) == [[1, 2, 3, 4, 5]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90d34107",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class GroupMetadata:\n",
    "    def __init__(self, num_partitions: int):\n",
    "        \"\"\"\n",
    "        Initialize a GroupMetadata object.\n",
    "\n",
    "        Args:\n",
    "            num_partitions: The number of partitions in the group.\n",
    "        \"\"\"\n",
    "        self.num_partitions = num_partitions\n",
    "        self.partitions_offsets: Dict[int, int] = {}\n",
    "        self.consumer_ids: List[uuid.UUID] = list()\n",
    "        self.partition_assignments: Dict[uuid.UUID, List[int]] = {}\n",
    "\n",
    "    def subscribe(self, consumer_id: uuid.UUID) -> None:\n",
    "        \"\"\"\n",
    "        Subscribe a consumer to the group.\n",
    "\n",
    "        Args:\n",
    "            consumer_id: The ID of the consumer.\n",
    "        \"\"\"\n",
    "        self.consumer_ids.append(consumer_id)\n",
    "        self.rebalance()\n",
    "\n",
    "    def unsubscribe(self, consumer_id: uuid.UUID) -> None:\n",
    "        \"\"\"\n",
    "        Unsubscribe a consumer from the group.\n",
    "\n",
    "        Args:\n",
    "            consumer_id: The ID of the consumer.\n",
    "        \"\"\"\n",
    "        self.consumer_ids.remove(consumer_id)\n",
    "        self.rebalance()\n",
    "\n",
    "    def rebalance(self) -> None:\n",
    "        \"\"\"\n",
    "        Rebalance the group's partition assignments.\n",
    "        \"\"\"\n",
    "        if len(self.consumer_ids) == 0:\n",
    "            self.partition_assignments = {}\n",
    "        else:\n",
    "            partitions_per_actor = self.num_partitions // len(self.consumer_ids)\n",
    "            if self.num_partitions % len(self.consumer_ids) != 0:\n",
    "                partitions_per_actor += 1\n",
    "            self.assign_partitions(partitions_per_actor)\n",
    "\n",
    "    def assign_partitions(self, partitions_per_actor: int) -> None:\n",
    "        partitions = [i for i in range(self.num_partitions)]\n",
    "\n",
    "        partitions_split = split_list(partitions, partitions_per_actor)\n",
    "        self.partition_assignments = {\n",
    "            self.consumer_ids[i]: partition_split\n",
    "            for i, partition_split in enumerate(partitions_split)\n",
    "        }\n",
    "\n",
    "    def get_partitions(\n",
    "        self, consumer_id: uuid.UUID\n",
    "    ) -> Tuple[List[int], Dict[int, Optional[int]]]:\n",
    "        \"\"\"\n",
    "        Get the partition assignments and offsets for a consumer.\n",
    "\n",
    "        Args:\n",
    "            consumer_id: The ID of the consumer.\n",
    "\n",
    "        Returns:\n",
    "            A tuple containing the partition assignments and offsets.\n",
    "        \"\"\"\n",
    "        partition_assignments = self.partition_assignments.get(consumer_id, [])\n",
    "        partition_offsets_assignments = {\n",
    "            partition: self.partitions_offsets.get(partition, None)\n",
    "            for partition in partition_assignments\n",
    "        }\n",
    "        return partition_assignments, partition_offsets_assignments\n",
    "\n",
    "    def set_offset(self, partition: int, offset: int) -> None:\n",
    "        \"\"\"\n",
    "        Set the offset for a partition.\n",
    "\n",
    "        Args:\n",
    "            partition: The partition index.\n",
    "            offset: The offset to set.\n",
    "        \"\"\"\n",
    "        self.partitions_offsets[partition] = offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a230027",
   "metadata": {},
   "outputs": [],
   "source": [
    "group_meta = GroupMetadata(num_partitions=3)\n",
    "\n",
    "# subscribe first consumer\n",
    "consumer_id_1 = uuid.uuid4()\n",
    "group_meta.subscribe(consumer_id_1)\n",
    "# check partitions\n",
    "assert group_meta.get_partitions(consumer_id_1)[0] == [0, 1, 2]\n",
    "\n",
    "# subscribe second consumer\n",
    "consumer_id_2 = uuid.uuid4()\n",
    "group_meta.subscribe(consumer_id_2)\n",
    "# check partitions\n",
    "assert group_meta.get_partitions(consumer_id_1)[0] == [0, 1]\n",
    "assert group_meta.get_partitions(consumer_id_2)[0] == [2]\n",
    "\n",
    "# subscribe third consumer\n",
    "consumer_id_3 = uuid.uuid4()\n",
    "group_meta.subscribe(consumer_id_3)\n",
    "# check partitions\n",
    "assert group_meta.get_partitions(consumer_id_1)[0] == [0]\n",
    "assert group_meta.get_partitions(consumer_id_2)[0] == [1]\n",
    "assert group_meta.get_partitions(consumer_id_3)[0] == [2]\n",
    "\n",
    "# subscribe fourth consumer\n",
    "# subscribe third consumer\n",
    "consumer_id_4 = uuid.uuid4()\n",
    "group_meta.subscribe(consumer_id_4)\n",
    "# check partitions\n",
    "assert group_meta.get_partitions(consumer_id_1)[0] == [0]\n",
    "assert group_meta.get_partitions(consumer_id_2)[0] == [1]\n",
    "assert group_meta.get_partitions(consumer_id_3)[0] == [2]\n",
    "assert group_meta.get_partitions(consumer_id_4)[0] == []  # fourth consumer is starving\n",
    "\n",
    "# Unsubscribe one consumer\n",
    "group_meta.unsubscribe(consumer_id_3)\n",
    "# check partitions\n",
    "assert group_meta.get_partitions(consumer_id_1)[0] == [0]\n",
    "assert group_meta.get_partitions(consumer_id_2)[0] == [1]\n",
    "assert group_meta.get_partitions(consumer_id_4)[0] == [2], group_meta.get_partitions(\n",
    "    consumer_id_4\n",
    ")\n",
    "\n",
    "# Unsubscribe all but one consumer\n",
    "group_meta.unsubscribe(consumer_id_1)\n",
    "group_meta.unsubscribe(consumer_id_4)\n",
    "assert group_meta.get_partitions(consumer_id_2)[0] == [0, 1, 2]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e25de69",
   "metadata": {},
   "source": [
    "## Kafka broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46e96510",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@classcontextmanager()\n",
    "class InMemoryBroker:\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_partitions: int = 1,\n",
    "    ):\n",
    "        self.num_partitions = num_partitions\n",
    "        self.topics: Dict[Tuple[str, str], KafkaTopic] = {}\n",
    "        self.topic_groups: Dict[Tuple[str, str, str], GroupMetadata] = {}\n",
    "        self.is_started: bool = False\n",
    "\n",
    "    def connect(self) -> uuid.UUID:\n",
    "        return uuid.uuid4()\n",
    "\n",
    "    def dissconnect(self, consumer_id: uuid.UUID) -> None:\n",
    "        \"\"\"\n",
    "        Disconnect a consumer from the broker.\n",
    "\n",
    "        Args:\n",
    "            consumer_id: The ID of the consumer.\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "    def subscribe(\n",
    "        self, bootstrap_server: str, topic: str, group: str, consumer_id: uuid.UUID\n",
    "    ) -> None:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def unsubscribe(\n",
    "        self, bootstrap_server: str, topic: str, group: str, consumer_id: uuid.UUID\n",
    "    ) -> None:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def read(  # type: ignore\n",
    "        self,\n",
    "        *,\n",
    "        bootstrap_server: str,\n",
    "        topic: str,\n",
    "        group: str,\n",
    "        consumer_id: uuid.UUID,\n",
    "        auto_offset_reset: str,\n",
    "    ) -> Dict[TopicPartition, List[KafkaRecord]]:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    def write(  # type: ignore\n",
    "        self,\n",
    "        *,\n",
    "        bootstrap_server: str,\n",
    "        topic: str,\n",
    "        value: bytes,\n",
    "        key: Optional[bytes] = None,\n",
    "        partition: Optional[int] = None,\n",
    "    ) -> RecordMetadata:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @contextmanager\n",
    "    def lifecycle(self) -> Iterator[\"InMemoryBroker\"]:\n",
    "        \"\"\"\n",
    "        Context manager for the lifecycle of the in-memory broker.\n",
    "\n",
    "        Yields:\n",
    "            An instance of the in-memory broker.\n",
    "        \"\"\"\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    async def _start(self) -> str:\n",
    "        \"\"\"\n",
    "        Start the in-memory broker.\n",
    "\n",
    "        Returns:\n",
    "            The address of the broker.\n",
    "        \"\"\"\n",
    "        logger.info(\"InMemoryBroker._start() called\")\n",
    "        self.__enter__()  # type: ignore\n",
    "        return \"localbroker:0\"\n",
    "\n",
    "    async def _stop(self) -> None:\n",
    "        \"\"\"\n",
    "        Stop the in-memory broker.\n",
    "        \"\"\"\n",
    "        logger.info(\"InMemoryBroker._stop() called\")\n",
    "        self.__exit__(None, None, None)  # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4a4bcae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def subscribe(\n",
    "    self: InMemoryBroker,\n",
    "    bootstrap_server: str,\n",
    "    topic: str,\n",
    "    group: str,\n",
    "    consumer_id: uuid.UUID,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Subscribe a consumer to a topic group.\n",
    "\n",
    "    Args:\n",
    "        bootstrap_server: The bootstrap server address.\n",
    "        topic: The topic to subscribe to.\n",
    "        group: The group to join.\n",
    "        consumer_id: The ID of the consumer.\n",
    "    \"\"\"\n",
    "    if (bootstrap_server, topic) not in self.topics:\n",
    "        self.topics[(bootstrap_server, topic)] = KafkaTopic(\n",
    "            topic=topic, num_partitions=self.num_partitions\n",
    "        )\n",
    "\n",
    "    group_meta = self.topic_groups.get(\n",
    "        (bootstrap_server, topic, group), GroupMetadata(self.num_partitions)\n",
    "    )\n",
    "    group_meta.subscribe(consumer_id)\n",
    "    self.topic_groups[(bootstrap_server, topic, group)] = group_meta\n",
    "\n",
    "\n",
    "@patch\n",
    "def unsubscribe(\n",
    "    self: InMemoryBroker,\n",
    "    bootstrap_server: str,\n",
    "    topic: str,\n",
    "    group: str,\n",
    "    consumer_id: uuid.UUID,\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Unsubscribe a consumer from a topic group.\n",
    "\n",
    "    Args:\n",
    "        bootstrap_server: The bootstrap server address.\n",
    "        topic: The topic to unsubscribe from.\n",
    "        group: The group to leave.\n",
    "        consumer_id: The ID of the consumer.\n",
    "    \"\"\"\n",
    "    self.topic_groups[(bootstrap_server, topic, group)].unsubscribe(consumer_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d4ac05b",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic = \"topic1\"\n",
    "bootstrap_server = \"localhost:9092\"\n",
    "consumer_group = \"my_group\"\n",
    "\n",
    "broker = InMemoryBroker()\n",
    "\n",
    "with pytest.raises(KeyError):\n",
    "    broker.topic_groups[(bootstrap_server, topic, consumer_group)]\n",
    "\n",
    "consumer_id = broker.connect()\n",
    "\n",
    "broker.subscribe(bootstrap_server, topic, consumer_group, consumer_id)\n",
    "broker.topic_groups[(bootstrap_server, topic, consumer_group)]\n",
    "\n",
    "broker.unsubscribe(bootstrap_server, topic, consumer_group, consumer_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72375d67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def write(  # type: ignore\n",
    "    self: InMemoryBroker,\n",
    "    *,\n",
    "    bootstrap_server: str,\n",
    "    topic: str,\n",
    "    value: bytes,\n",
    "    key: Optional[bytes] = None,\n",
    "    partition: Optional[int] = None,\n",
    ") -> RecordMetadata:\n",
    "    \"\"\"\n",
    "    Write a message to a topic.\n",
    "\n",
    "    Args:\n",
    "        bootstrap_server: The bootstrap server address.\n",
    "        topic: The topic to write the message to.\n",
    "        value: The value of the message.\n",
    "        key: The key associated with the message.\n",
    "        partition: The partition ID to write the message to.\n",
    "\n",
    "    Returns:\n",
    "        The metadata of the written message.\n",
    "    \"\"\"\n",
    "    if (bootstrap_server, topic) not in self.topics:\n",
    "        self.topics[(bootstrap_server, topic)] = KafkaTopic(\n",
    "            topic=topic, num_partitions=self.num_partitions\n",
    "        )\n",
    "\n",
    "    return self.topics[(bootstrap_server, topic)].write(\n",
    "        value, key=key, partition=partition\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68339dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "for key in [None, b\"some_key\"]:\n",
    "    topic = \"my_topic\"\n",
    "    bootstrap_server = \"localhost:9092\"\n",
    "    value = b\"msg\"\n",
    "\n",
    "    broker = InMemoryBroker(num_partitions=3)\n",
    "\n",
    "    record_meta = broker.write(\n",
    "        bootstrap_server=bootstrap_server, topic=topic, value=value, key=key\n",
    "    )\n",
    "\n",
    "    assert record_meta.topic == \"my_topic\"\n",
    "    assert record_meta.offset == 0\n",
    "\n",
    "    expected_msgs = [\n",
    "        KafkaRecord(\n",
    "            topic=\"my_topic\",\n",
    "            partition=record_meta.partition,\n",
    "            key=key,\n",
    "            value=b\"msg\",\n",
    "            offset=0,\n",
    "        )\n",
    "    ]\n",
    "\n",
    "    topic_partition, actual_msgs, new_offset = broker.topics[\n",
    "        (bootstrap_server, topic)\n",
    "    ].read(partition=record_meta.partition, offset=record_meta.offset)\n",
    "\n",
    "    assert actual_msgs == expected_msgs\n",
    "    assert topic_partition == record_meta.topic_partition\n",
    "    assert new_offset == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f44154f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "def read(  # type: ignore\n",
    "    self: InMemoryBroker,\n",
    "    *,\n",
    "    bootstrap_server: str,\n",
    "    topic: str,\n",
    "    group: str,\n",
    "    consumer_id: uuid.UUID,\n",
    "    auto_offset_reset: str,\n",
    ") -> Dict[TopicPartition, List[KafkaRecord]]:\n",
    "    \"\"\"\n",
    "    Read messages from a topic group.\n",
    "\n",
    "    Args:\n",
    "        bootstrap_server: The bootstrap server address.\n",
    "        topic: The topic to read messages from.\n",
    "        group: The group to read messages for.\n",
    "        consumer_id: The ID of the consumer.\n",
    "        auto_offset_reset: The strategy to use when the consumer does not have a valid offset for the group.\n",
    "\n",
    "    Returns:\n",
    "        A dictionary containing the messages retrieved from each topic partition.\n",
    "    \"\"\"\n",
    "    group_meta = self.topic_groups[(bootstrap_server, topic, group)]\n",
    "    partitions, offsets = group_meta.get_partitions(consumer_id)\n",
    "\n",
    "    if len(partitions) == 0:\n",
    "        return {}\n",
    "\n",
    "    partitions_data = {}\n",
    "\n",
    "    for partition in partitions:\n",
    "        offset = offsets[partition]\n",
    "\n",
    "        if offset is None:\n",
    "            offset = (\n",
    "                self.topics[(bootstrap_server, topic)].latest_offset(partition)\n",
    "                if auto_offset_reset == \"latest\"\n",
    "                else 0\n",
    "            )\n",
    "\n",
    "        topic_partition, data, offset = self.topics[(bootstrap_server, topic)].read(\n",
    "            partition, offset\n",
    "        )\n",
    "\n",
    "        partitions_data[topic_partition] = data\n",
    "        group_meta.set_offset(partition, offset)\n",
    "\n",
    "    return partitions_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f2edf62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check subscribing and reading from empty partitions for same group\n",
    "\n",
    "topic = \"topic1\"\n",
    "bootstrap_server = \"localhost:9092\"\n",
    "consumer_group = \"my_group\"\n",
    "\n",
    "broker = InMemoryBroker(num_partitions=3)\n",
    "\n",
    "consumer_id_1 = broker.connect()\n",
    "broker.subscribe(bootstrap_server, topic, consumer_group, consumer_id_1)\n",
    "\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    group=consumer_group,\n",
    "    consumer_id=consumer_id_1,\n",
    "    auto_offset_reset=\"latest\",\n",
    ") == {\n",
    "    TopicPartition(topic=topic, partition=0): [],\n",
    "    TopicPartition(topic=topic, partition=1): [],\n",
    "    TopicPartition(topic=topic, partition=2): [],\n",
    "}\n",
    "\n",
    "consumer_id_2 = broker.connect()\n",
    "broker.subscribe(bootstrap_server, topic, consumer_group, consumer_id_2)\n",
    "\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    group=consumer_group,\n",
    "    consumer_id=consumer_id_1,\n",
    "    auto_offset_reset=\"latest\",\n",
    ") == {\n",
    "    TopicPartition(topic=topic, partition=0): [],\n",
    "    TopicPartition(topic=topic, partition=1): [],\n",
    "}\n",
    "\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    group=consumer_group,\n",
    "    consumer_id=consumer_id_2,\n",
    "    auto_offset_reset=\"latest\",\n",
    ") == {\n",
    "    TopicPartition(topic=topic, partition=2): [],\n",
    "}\n",
    "\n",
    "broker.unsubscribe(bootstrap_server, topic, consumer_group, consumer_id_1)\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    group=consumer_group,\n",
    "    consumer_id=consumer_id_2,\n",
    "    auto_offset_reset=\"latest\",\n",
    ") == {\n",
    "    TopicPartition(topic=topic, partition=0): [],\n",
    "    TopicPartition(topic=topic, partition=1): [],\n",
    "    TopicPartition(topic=topic, partition=2): [],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a30de58e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# check writing to partitions\n",
    "\n",
    "topic = \"topic1\"\n",
    "bootstrap_server = \"localhost:9092\"\n",
    "consumer_group = \"my_group\"\n",
    "\n",
    "broker = InMemoryBroker(num_partitions=1)\n",
    "\n",
    "consumer_id_1 = broker.connect()\n",
    "broker.subscribe(bootstrap_server, topic, consumer_group, consumer_id_1)\n",
    "\n",
    "record_meta = broker.write(bootstrap_server=bootstrap_server, topic=topic, value=b\"msg\")\n",
    "\n",
    "assert record_meta.topic == topic\n",
    "assert record_meta.partition == 0\n",
    "assert record_meta.topic_partition == TopicPartition(topic=topic, partition=0)\n",
    "assert record_meta.offset == 0\n",
    "\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    consumer_id=consumer_id_1,\n",
    "    group=consumer_group,\n",
    "    auto_offset_reset=\"earliest\",\n",
    ") == {\n",
    "    TopicPartition(topic=topic, partition=0): [\n",
    "        KafkaRecord(topic=topic, partition=0, key=None, value=b\"msg\", offset=0)\n",
    "    ]\n",
    "}\n",
    "\n",
    "broker.write(bootstrap_server=bootstrap_server, topic=topic, value=b\"msg\")\n",
    "\n",
    "consumer_group_new = \"another_group\"\n",
    "\n",
    "consumer_id_2 = broker.connect()\n",
    "broker.subscribe(bootstrap_server, topic, consumer_group_new, consumer_id_2)\n",
    "\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    consumer_id=consumer_id_2,\n",
    "    group=consumer_group_new,\n",
    "    auto_offset_reset=\"latest\",\n",
    ") == {TopicPartition(topic=topic, partition=0): []}\n",
    "\n",
    "assert broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    consumer_id=consumer_id_1,\n",
    "    group=consumer_group,\n",
    "    auto_offset_reset=\"latest\",\n",
    ") == {\n",
    "    TopicPartition(topic=topic, partition=0): [\n",
    "        KafkaRecord(topic=topic, partition=0, key=None, value=b\"msg\", offset=1)\n",
    "    ]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c59295ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic = \"my_topic\"\n",
    "bootstrap_server = \"localhost:9092\"\n",
    "group = \"my_group\"\n",
    "\n",
    "in_memory_broker = InMemoryBroker()\n",
    "\n",
    "consumer_id = in_memory_broker.connect()\n",
    "\n",
    "with pytest.raises(KeyError) as e:\n",
    "    in_memory_broker.read(\n",
    "        bootstrap_server=bootstrap_server,\n",
    "        topic=topic,\n",
    "        group=group,\n",
    "        consumer_id=consumer_id,\n",
    "        auto_offset_reset=\"latest\",\n",
    "    )\n",
    "\n",
    "in_memory_broker.subscribe(\n",
    "    bootstrap_server=bootstrap_server, topic=topic, group=group, consumer_id=consumer_id\n",
    ")\n",
    "\n",
    "msg = in_memory_broker.read(\n",
    "    bootstrap_server=bootstrap_server,\n",
    "    topic=topic,\n",
    "    group=group,\n",
    "    consumer_id=consumer_id,\n",
    "    auto_offset_reset=\"earliest\",\n",
    ")\n",
    "assert msg == {TopicPartition(topic=topic, partition=0): []}, msg"
   ]
  },
  {
   "cell_type": "raw",
   "id": "cee0ab7e",
   "metadata": {},
   "source": [
    "#| notest\n",
    "\n",
    "with ApacheKafkaBroker([\"my_topic\"], apply_nest_asyncio=True) as bootstrap_servers:\n",
    "    producer = AIOKafkaProducer(bootstrap_servers=bootstrap_servers)\n",
    "    await producer.start()\n",
    "    for _ in range(1000):\n",
    "        record = await producer.send(topic=\"not_my_topic\", value=b\"not my message\")\n",
    "    await producer.stop()"
   ]
  },
  {
   "cell_type": "raw",
   "id": "fe028489",
   "metadata": {},
   "source": [
    "#| notest\n",
    "\n",
    "with ApacheKafkaBroker([\"my_topic\"], apply_nest_asyncio=True) as bootstrap_servers:\n",
    "    consumer = AIOKafkaConsumer(\"my_topic\", bootstrap_servers=bootstrap_servers)\n",
    "    await consumer.start()\n",
    "    print(\"getmany()...\")\n",
    "    msg = await consumer.getmany(timeout_ms=0)\n",
    "    print(\"exiting...\")\n",
    "    await consumer.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c5b3a9",
   "metadata": {},
   "source": [
    "## Consumer patching\n",
    "\n",
    "We need to patch AIOKafkaConsumer methods so that we can redirect the consumer to our local kafka broker.\n",
    "\n",
    "Patched methods:\n",
    "\n",
    "- [x] \\_\\_init\\_\\_\n",
    "- [x] start\n",
    "- [x] subscribe\n",
    "- [x] stop\n",
    "- [x] getmany"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c85ca42b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "# InMemoryConsumer\n",
    "class InMemoryConsumer:\n",
    "    def __init__(\n",
    "        self,\n",
    "        broker: InMemoryBroker,\n",
    "    ) -> None:\n",
    "        self.broker = broker\n",
    "        self._id: Optional[uuid.UUID] = None\n",
    "        self._auto_offset_reset: str = \"latest\"\n",
    "        self._group_id: Optional[str] = None\n",
    "        self._topics: List[str] = list()\n",
    "        self._bootstrap_servers = \"\"\n",
    "\n",
    "    @delegates(AIOKafkaConsumer)\n",
    "    def __call__(self, **kwargs: Any) -> \"InMemoryConsumer\":\n",
    "        defaults = _get_default_kwargs_from_sig(InMemoryConsumer.__call__, **kwargs)\n",
    "        consume_copy = InMemoryConsumer(self.broker)\n",
    "        consume_copy._auto_offset_reset = defaults[\"auto_offset_reset\"]\n",
    "        consume_copy._bootstrap_servers = (\n",
    "            \"\".join(defaults[\"bootstrap_servers\"])\n",
    "            if isinstance(defaults[\"bootstrap_servers\"], list)\n",
    "            else defaults[\"bootstrap_servers\"]\n",
    "        )\n",
    "\n",
    "        consume_copy._group_id = (\n",
    "            defaults[\"group_id\"]\n",
    "            if defaults[\"group_id\"] is not None\n",
    "            else \"\".join(random.choices(string.ascii_letters, k=10))  # nosec\n",
    "        )\n",
    "        return consume_copy\n",
    "\n",
    "    @delegates(AIOKafkaConsumer.start)\n",
    "    async def start(self, **kwargs: Any) -> None:\n",
    "        pass\n",
    "\n",
    "    @delegates(AIOKafkaConsumer.stop)\n",
    "    async def stop(self, **kwargs: Any) -> None:\n",
    "        pass\n",
    "\n",
    "    @delegates(AIOKafkaConsumer.subscribe)\n",
    "    def subscribe(self, topics: List[str], **kwargs: Any) -> None:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaConsumer.getmany)\n",
    "    async def getmany(  # type: ignore\n",
    "        self, **kwargs: Any\n",
    "    ) -> Dict[TopicPartition, List[ConsumerRecord]]:\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "677208e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:22.629 [ERROR] asyncio: Unclosed AIOKafkaConsumer\n",
      "consumer: <aiokafka.consumer.consumer.AIOKafkaConsumer object>\n",
      "23-06-26 08:04:22.631 [ERROR] asyncio: Unclosed AIOKafkaConsumer\n",
      "consumer: <aiokafka.consumer.consumer.AIOKafkaConsumer object>\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "\n",
    "for cls in [ConsumerClass, AIOKafkaConsumer]:\n",
    "    consumer = cls()\n",
    "    assert consumer._auto_offset_reset == \"latest\"\n",
    "\n",
    "    consumer = cls(auto_offset_reset=\"earliest\")\n",
    "    assert consumer._auto_offset_reset == \"earliest\", consumer._auto_offset_reset\n",
    "\n",
    "    consumer = cls(auto_offset_reset=\"whatever\")\n",
    "    assert consumer._auto_offset_reset == \"whatever\"\n",
    "\n",
    "    await consumer.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28e4df06",
   "metadata": {},
   "source": [
    "Patching start so that we don't try to start the real AIOKafkaConsumer instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f53f9da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaConsumer.start)\n",
    "async def start(self: InMemoryConsumer, **kwargs: Any) -> None:\n",
    "    \"\"\"\n",
    "    Start consuming messages from the connected broker.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has already been called without calling stop() first.\n",
    "    \"\"\"\n",
    "    logger.info(\"AIOKafkaConsumer patched start() called()\")\n",
    "    if self._id is not None:\n",
    "        raise RuntimeError(\n",
    "            \"Consumer start() already called! Run consumer stop() before running start() again\"\n",
    "        )\n",
    "    self._id = self.broker.connect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "072dd1f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:33.493 [INFO] __main__: AIOKafkaConsumer patched start() called()\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "\n",
    "for cls in [ConsumerClass]:\n",
    "    consumer = cls()\n",
    "    await consumer.start()\n",
    "    await consumer.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21c11cd3",
   "metadata": {},
   "source": [
    "Patching subscribe so that we can connect to our Local, in-memory, Kafka broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1de72311",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch  # type: ignore\n",
    "@delegates(AIOKafkaConsumer.subscribe)\n",
    "def subscribe(self: InMemoryConsumer, topics: List[str], **kwargs: Any) -> None:\n",
    "    \"\"\"\n",
    "    Subscribe to a list of topics for consuming messages.\n",
    "\n",
    "    Args:\n",
    "        topics: A list of topics to subscribe to.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has not been called before calling subscribe().\n",
    "    \"\"\"\n",
    "    logger.info(\"AIOKafkaConsumer patched subscribe() called\")\n",
    "    if self._id is None:\n",
    "        raise RuntimeError(\"Consumer start() not called! Run consumer start() first\")\n",
    "    logger.info(f\"AIOKafkaConsumer.subscribe(), subscribing to: {topics}\")\n",
    "    for topic in topics:\n",
    "        self.broker.subscribe(\n",
    "            bootstrap_server=self._bootstrap_servers,\n",
    "            consumer_id=self._id,\n",
    "            topic=topic,\n",
    "            group=self._group_id,  # type: ignore\n",
    "        )\n",
    "        self._topics.append(topic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "945e08d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:37.008 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:04:37.010 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:04:37.010 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['my_topic']\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "consumer = ConsumerClass()\n",
    "await consumer.start()\n",
    "consumer.subscribe([\"my_topic\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd80a4d5",
   "metadata": {},
   "source": [
    "Patching stop so that be dont break anything by calling the real AIOKafkaConsumer stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc82405e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaConsumer.stop)\n",
    "async def stop(self: InMemoryConsumer, **kwargs: Any) -> None:\n",
    "    \"\"\"\n",
    "    Stop consuming messages from the connected broker.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has not been called before calling stop().\n",
    "    \"\"\"\n",
    "    logger.info(\"AIOKafkaConsumer patched stop() called\")\n",
    "    if self._id is None:\n",
    "        raise RuntimeError(\"Consumer start() not called! Run consumer start() first\")\n",
    "    for topic in self._topics:\n",
    "        self.broker.unsubscribe(\n",
    "            bootstrap_server=self._bootstrap_servers,\n",
    "            topic=topic,\n",
    "            group=self._group_id,  # type: ignore\n",
    "            consumer_id=self._id,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc667214",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:38.785 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:04:38.786 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:04:38.787 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['my_topic']\n",
      "23-06-26 08:04:38.787 [INFO] __main__: AIOKafkaConsumer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "consumer = ConsumerClass()\n",
    "\n",
    "await consumer.start()\n",
    "consumer.subscribe([\"my_topic\"])\n",
    "await consumer.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0c733b4",
   "metadata": {},
   "source": [
    "Patching getmany so that the messages are pulled from our Local, in-memory, Kafka broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0663c9b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaConsumer.getmany)\n",
    "async def getmany(  # type: ignore\n",
    "    self: InMemoryConsumer, **kwargs: Any\n",
    ") -> Dict[TopicPartition, List[ConsumerRecord]]:\n",
    "    \"\"\"\n",
    "    Retrieve messages from the subscribed topics.\n",
    "\n",
    "    Returns:\n",
    "        A dictionary containing the retrieved messages from each topic partition.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has not been called before calling getmany().\n",
    "    \"\"\"\n",
    "    await asyncio.sleep(0)\n",
    "    for topic in self._topics:\n",
    "        return self.broker.read(\n",
    "            bootstrap_server=self._bootstrap_servers,\n",
    "            topic=topic,\n",
    "            consumer_id=self._id,  # type: ignore\n",
    "            group=self._group_id,  # type: ignore\n",
    "            auto_offset_reset=self._auto_offset_reset,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4940bac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:40.641 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:04:40.644 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:04:40.646 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['my_topic']\n",
      "23-06-26 08:04:40.649 [INFO] __main__: AIOKafkaConsumer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "consumer = ConsumerClass(auto_offset_reset=\"latest\")\n",
    "\n",
    "await consumer.start()\n",
    "\n",
    "consumer.subscribe([\"my_topic\"])\n",
    "await consumer.getmany()\n",
    "\n",
    "await consumer.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "723468f0",
   "metadata": {},
   "source": [
    "## Producer patching\n",
    "\n",
    "We need to patch AIOKafkaProducer methods so that we can redirect the producer to our local kafka broker\n",
    "\n",
    "- [x] \\_\\_init\\_\\_\n",
    "- [x] start\n",
    "- [x] stop\n",
    "- [x] send"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad3d6b0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class InMemoryProducer:\n",
    "    def __init__(self, broker: InMemoryBroker, **kwargs: Any) -> None:\n",
    "        self.broker = broker\n",
    "        self.id: Optional[uuid.UUID] = None\n",
    "        self._bootstrap_servers = \"\"\n",
    "\n",
    "    @delegates(AIOKafkaProducer)\n",
    "    def __call__(self, **kwargs: Any) -> \"InMemoryProducer\":\n",
    "        defaults = _get_default_kwargs_from_sig(InMemoryConsumer.__call__, **kwargs)\n",
    "        producer_copy = InMemoryProducer(self.broker)\n",
    "        producer_copy._bootstrap_servers = (\n",
    "            \"\".join(defaults[\"bootstrap_servers\"])\n",
    "            if isinstance(defaults[\"bootstrap_servers\"], list)\n",
    "            else defaults[\"bootstrap_servers\"]\n",
    "        )\n",
    "        return producer_copy\n",
    "\n",
    "    @delegates(AIOKafkaProducer.start)\n",
    "    async def start(self, **kwargs: Any) -> None:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaProducer.stop)\n",
    "    async def stop(self, **kwargs: Any) -> None:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaProducer.send)\n",
    "    async def send(  # type: ignore\n",
    "        self,\n",
    "        topic: str,\n",
    "        msg: bytes,\n",
    "        key: Optional[bytes] = None,\n",
    "        **kwargs: Any,\n",
    "    ):\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaProducer.partitions_for)\n",
    "    async def partitions_for(self, topic: str) -> List[int]:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaProducer._partition)\n",
    "    def _partition(\n",
    "        self, topic: str, arg1: Any, arg2: Any, arg3: Any, key: bytes, arg4: Any\n",
    "    ) -> int:\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaProducer.create_batch)\n",
    "    def create_batch(self) -> \"MockBatch\":\n",
    "        raise NotImplementedError()\n",
    "\n",
    "    @delegates(AIOKafkaProducer.send_batch)\n",
    "    async def send_batch(self, batch: \"MockBatch\", topic: str, partition: Any) -> None:\n",
    "        raise NotImplementedError()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e509e2a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "producer_cls = InMemoryProducer(None)\n",
    "\n",
    "producer = producer_cls()\n",
    "assert producer._bootstrap_servers == \"localhost\"\n",
    "\n",
    "producer = producer_cls(bootstrap_servers=\"kafka.airt.ai\")\n",
    "assert producer._bootstrap_servers == \"kafka.airt.ai\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bedf322a",
   "metadata": {},
   "source": [
    "Patching AIOKafkaProducer start so that we mock the startup procedure of AIOKafkaProducer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "488ac5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch  # type: ignore\n",
    "@delegates(AIOKafkaProducer.start)\n",
    "async def start(self: InMemoryProducer, **kwargs: Any) -> None:\n",
    "    \"\"\"\n",
    "    Start the in-memory producer.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has already been called without calling stop() first.\n",
    "    \"\"\"\n",
    "    logger.info(\"AIOKafkaProducer patched start() called()\")\n",
    "    if self.id is not None:\n",
    "        raise RuntimeError(\n",
    "            \"Producer start() already called! Run producer stop() before running start() again\"\n",
    "        )\n",
    "    self.id = self.broker.connect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f250c614",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:43.984 [INFO] __main__: AIOKafkaProducer patched start() called()\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "await producer.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f9e9be9",
   "metadata": {},
   "source": [
    "Patching AIOKafkaProducerStop so that we don't uniintentionally try to stop a real instance of AIOKafkaProducer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32412969",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch  # type: ignore\n",
    "@delegates(AIOKafkaProducer.stop)\n",
    "async def stop(self: InMemoryProducer, **kwargs: Any) -> None:\n",
    "    \"\"\"\n",
    "    Stop the in-memory producer.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has not been called before calling stop().\n",
    "    \"\"\"\n",
    "    logger.info(\"AIOKafkaProducer patched stop() called\")\n",
    "    if self.id is None:\n",
    "        raise RuntimeError(\"Producer start() not called! Run producer start() first\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f4a1fbb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:04:47.703 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:04:47.704 [INFO] __main__: AIOKafkaProducer patched stop() called\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "await producer.start()\n",
    "await producer.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c77a56",
   "metadata": {},
   "source": [
    "Patching AIOKafkaProducer send so that we redirect sent messages to Local, in-memory, Kafka broker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f42a03e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaProducer.send)\n",
    "async def send(  # type: ignore\n",
    "    self: InMemoryProducer,\n",
    "    topic: str,\n",
    "    msg: bytes,\n",
    "    key: Optional[bytes] = None,\n",
    "    partition: Optional[int] = None,\n",
    "    **kwargs: Any,\n",
    "):  # asyncio.Task[RecordMetadata]\n",
    "    \"\"\"\n",
    "    Send a message to the specified topic.\n",
    "\n",
    "    Args:\n",
    "        topic: The topic to send the message to.\n",
    "        msg: The message to send.\n",
    "        key: The key associated with the message (optional).\n",
    "        partition: The partition to send the message to (optional).\n",
    "        **kwargs: Additional arguments to be passed to AIOKafkaProducer.send().\n",
    "\n",
    "    Returns:\n",
    "        A task that resolves to the RecordMetadata of the sent message.\n",
    "\n",
    "    Raises:\n",
    "        RuntimeError: If start() has not been called before calling send().\n",
    "    \"\"\"\n",
    "    if self.id is None:\n",
    "        raise RuntimeError(\"Producer start() not called! Run producer start() first\")\n",
    "\n",
    "    record = self.broker.write(\n",
    "        bootstrap_server=self._bootstrap_servers,\n",
    "        topic=topic,\n",
    "        value=msg,\n",
    "        key=key,\n",
    "        partition=partition,\n",
    "    )\n",
    "\n",
    "    async def _f(record: ConsumerRecord = record) -> RecordMetadata:  # type: ignore\n",
    "        return record\n",
    "\n",
    "    return asyncio.create_task(_f())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fda1d9b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:05:09.942 [INFO] __main__: AIOKafkaProducer patched start() called()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RecordMetadata(topic='my_topic', partition=0, topic_partition=TopicPartition(topic='my_topic', partition=0), offset=0, timestamp=1680602752070, timestamp_type=0, log_start_offset=0)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "await producer.start()\n",
    "msg_fut = await producer.send(\"my_topic\", b\"some_msg\")\n",
    "await msg_fut"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eecd05b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaProducer.partitions_for)\n",
    "async def partitions_for(self: InMemoryProducer, topic: str) -> List[int]:\n",
    "    \"\"\"\n",
    "    Retrieve the list of partitions for the specified topic.\n",
    "\n",
    "    Args:\n",
    "        topic: The topic to get the partitions for.\n",
    "\n",
    "    Returns:\n",
    "        A list of partition IDs.\n",
    "    \"\"\"\n",
    "    return [i for i in range(self.broker.num_partitions)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea87579a",
   "metadata": {},
   "outputs": [],
   "source": [
    "broker = InMemoryBroker(num_partitions=5)\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "assert len(await producer.partitions_for(\"some_topic\")) == 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab634be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaProducer._partition)\n",
    "def _partition(\n",
    "    self: InMemoryProducer,\n",
    "    topic: str,\n",
    "    arg1: Any,\n",
    "    arg2: Any,\n",
    "    arg3: Any,\n",
    "    key: bytes,\n",
    "    arg4: Any,\n",
    ") -> int:\n",
    "    \"\"\"\n",
    "    Determine the partition to which the message should be sent.\n",
    "\n",
    "    Args:\n",
    "        topic: The topic to send the message to.\n",
    "        arg1, arg2, arg3, arg4: Additional arguments passed to the original AIOKafkaProducer._partition().\n",
    "\n",
    "    Returns:\n",
    "        The partition ID.\n",
    "    \"\"\"\n",
    "    return int(hashlib.sha256(key).hexdigest(), 16) % self.broker.num_partitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2038adc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "broker = InMemoryBroker(num_partitions=5)\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "partition = producer._partition(\"my_topic\", None, None, None, b\"key\", None)\n",
    "assert partition >= 0 and partition < 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4fcacc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "class MockBatch:\n",
    "    def __init__(self) -> None:\n",
    "        \"\"\"\n",
    "        Initialize an instance of MockBatch.\n",
    "        \"\"\"\n",
    "        self._batch: List[Tuple] = list()\n",
    "\n",
    "    def append(  # type: ignore\n",
    "        self, key: Optional[bytes], value: bytes, timestamp: int\n",
    "    ) -> RecordMetadata:\n",
    "        \"\"\"\n",
    "        Append a message to the batch.\n",
    "\n",
    "        Args:\n",
    "            key: The key associated with the message (optional).\n",
    "            value: The value of the message.\n",
    "            timestamp: The timestamp of the message.\n",
    "\n",
    "        Returns:\n",
    "            The RecordMetadata of the appended message.\n",
    "        \"\"\"\n",
    "        self._batch.append((key, value))\n",
    "        return RecordMetadata(\n",
    "            topic=\"\",\n",
    "            partition=0,\n",
    "            topic_partition=None,\n",
    "            offset=0,\n",
    "            timestamp=timestamp,\n",
    "            timestamp_type=0,\n",
    "            log_start_offset=0,\n",
    "        )\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaProducer.create_batch)\n",
    "def create_batch(self: InMemoryProducer) -> \"MockBatch\":\n",
    "    \"\"\"\n",
    "    Create a mock batch for the in-memory producer.\n",
    "\n",
    "    Returns:\n",
    "        A MockBatch instance.\n",
    "    \"\"\"\n",
    "    return MockBatch()\n",
    "\n",
    "\n",
    "@patch\n",
    "@delegates(AIOKafkaProducer.send_batch)\n",
    "async def send_batch(\n",
    "    self: InMemoryProducer, batch: \"MockBatch\", topic: str, partition: Any\n",
    ") -> None:\n",
    "    \"\"\"\n",
    "    Send a batch of messages to the specified topic and partition.\n",
    "\n",
    "    Args:\n",
    "        batch: The MockBatch containing the messages to send.\n",
    "        topic: The topic to send the batch of messages to.\n",
    "        partition: The partition to send the batch of messages to.\n",
    "    \"\"\"\n",
    "    for record in batch._batch:\n",
    "        self.broker.write(\n",
    "            bootstrap_server=self._bootstrap_servers,\n",
    "            topic=topic,\n",
    "            value=record[1],\n",
    "            key=record[0],\n",
    "            partition=partition,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62c9f00b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:05:14.471 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:05:14.472 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:05:14.473 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:05:14.473 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['my_topic']\n"
     ]
    }
   ],
   "source": [
    "bootstrap_server\n",
    "\n",
    "broker = InMemoryBroker()\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "await producer.start()\n",
    "\n",
    "batch = producer.create_batch()\n",
    "batch.append(b\"key\", b\"value\", 1)\n",
    "\n",
    "partition = producer._partition(\"my_topic\", None, None, None, b\"key\", None)\n",
    "await producer.send_batch(batch, topic, partition=partition)\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "consumer = ConsumerClass(auto_offset_reset=\"earliest\")\n",
    "\n",
    "await consumer.start()\n",
    "\n",
    "consumer.subscribe([\"my_topic\"])\n",
    "msgs = await consumer.getmany()\n",
    "assert len(msgs[TopicPartition(topic='my_topic', partition=0)]) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c37ae80",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:05:17.175 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:05:17.176 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:05:17.176 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:05:17.177 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['my_topic']\n"
     ]
    }
   ],
   "source": [
    "broker = InMemoryBroker()\n",
    "\n",
    "ProducerClass = InMemoryProducer(broker)\n",
    "producer = ProducerClass()\n",
    "\n",
    "await producer.start()\n",
    "\n",
    "batch = producer.create_batch()\n",
    "batch.append(b\"key\", b\"value\", 1)\n",
    "\n",
    "partitions = await producer.partitions_for(\"my_topic\")\n",
    "partition = random.choice(tuple(partitions))\n",
    "\n",
    "await producer.send_batch(batch, topic, partition=partition)\n",
    "\n",
    "ConsumerClass = InMemoryConsumer(broker)\n",
    "consumer = ConsumerClass(auto_offset_reset=\"earliest\")\n",
    "\n",
    "await consumer.start()\n",
    "\n",
    "consumer.subscribe([\"my_topic\"])\n",
    "msgs = await consumer.getmany()\n",
    "assert len(msgs[TopicPartition(topic='my_topic', partition=0)]) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e4d12b",
   "metadata": {},
   "source": [
    "## Add patching to InMemoryBroker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "446e0a24",
   "metadata": {},
   "outputs": [],
   "source": [
    "# | export\n",
    "\n",
    "\n",
    "@patch\n",
    "@contextmanager\n",
    "def lifecycle(self: InMemoryBroker) -> Iterator[InMemoryBroker]:\n",
    "    \"\"\"\n",
    "    Context manager for the lifecycle of the in-memory broker.\n",
    "\n",
    "    Yields:\n",
    "        An instance of the in-memory broker.\n",
    "    \"\"\"\n",
    "    logger.info(\n",
    "        \"InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\"\n",
    "    )\n",
    "    try:\n",
    "        logger.info(\"InMemoryBroker starting\")\n",
    "        \n",
    "        old_consumer = fastkafka._aiokafka_imports.AIOKafkaConsumer\n",
    "        old_producer = fastkafka._aiokafka_imports.AIOKafkaProducer\n",
    "        \n",
    "        fastkafka._aiokafka_imports.AIOKafkaConsumer = InMemoryConsumer(self)\n",
    "        fastkafka._aiokafka_imports.AIOKafkaProducer = InMemoryProducer(self)\n",
    "\n",
    "        self.is_started = True\n",
    "        yield self\n",
    "    finally:\n",
    "        logger.info(\"InMemoryBroker stopping\")\n",
    "\n",
    "        fastkafka._aiokafka_imports.AIOKafkaConsumer = old_consumer\n",
    "        fastkafka._aiokafka_imports.AIOKafkaProducer = old_producer\n",
    "\n",
    "        self.is_started = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc55ccb2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:05:22.344 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:05:22.345 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:05:22.345 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "assert fastkafka._aiokafka_imports.AIOKafkaConsumer == AIOKafkaConsumer\n",
    "assert fastkafka._aiokafka_imports.AIOKafkaProducer == AIOKafkaProducer\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    assert isinstance(fastkafka._aiokafka_imports.AIOKafkaConsumer, InMemoryConsumer)\n",
    "    assert isinstance(fastkafka._aiokafka_imports.AIOKafkaProducer, InMemoryProducer)\n",
    "    assert fastkafka._aiokafka_imports.AIOKafkaConsumer().broker == broker\n",
    "    assert fastkafka._aiokafka_imports.AIOKafkaProducer().broker == broker\n",
    "\n",
    "assert fastkafka._aiokafka_imports.AIOKafkaConsumer == AIOKafkaConsumer\n",
    "assert fastkafka._aiokafka_imports.AIOKafkaProducer == AIOKafkaProducer"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5335aea0",
   "metadata": {},
   "source": [
    "## Broker, consumer and producer integration tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4275bf97",
   "metadata": {},
   "outputs": [],
   "source": [
    "@asynccontextmanager\n",
    "async def create_consumer_and_producer(\n",
    "    auto_offset_reset: str = \"latest\",\n",
    ") -> AsyncIterator[Tuple[AIOKafkaConsumer, AIOKafkaProducer]]:\n",
    "    consumer = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "        auto_offset_reset=auto_offset_reset\n",
    "    )\n",
    "    producer = fastkafka._aiokafka_imports.AIOKafkaProducer()\n",
    "\n",
    "    await consumer.start()\n",
    "    await producer.start()\n",
    "\n",
    "    yield (consumer, producer)\n",
    "\n",
    "    await consumer.stop()\n",
    "    await producer.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a7688d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def checkEqual(L1, L2):\n",
    "    return len(L1) == len(L2) and sorted(L1) == sorted(L2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "800d6a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert checkEqual([1, 2], [3]) == False\n",
    "assert checkEqual([1, 2, 3], [3, 2, 1]) == True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebc21a6e",
   "metadata": {},
   "source": [
    "Sanity check, let's see if the messages are sent to broker and received by the consumer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f90249e0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:06:03.590 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:06:03.591 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:06:03.592 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:06:03.592 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:06:03.595 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:06:03.595 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:06:03.596 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:06:03.597 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:06:03.598 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "topic = \"test_topic\"\n",
    "sent_msgs = [f\"msg{i}\".encode(\"UTF-8\") for i in range(320)]\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    async with create_consumer_and_producer(auto_offset_reset=\"earliest\") as (\n",
    "        consumer,\n",
    "        producer,\n",
    "    ):\n",
    "        [await producer.send(topic, msg) for msg in sent_msgs]\n",
    "        consumer.subscribe([topic])\n",
    "        received = await consumer.getmany()\n",
    "        received_msgs = [msg.value for _, msgs in received.items() for msg in msgs]\n",
    "    assert checkEqual(\n",
    "        received_msgs, sent_msgs\n",
    "    ), f\"{sent_msgs=}\\n{received_msgs=}\\n{data=}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f40fa9ed",
   "metadata": {},
   "source": [
    "Check if only subscribed topic messages are received by the consumer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "839a6755",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:06:11.830 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:06:11.831 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:06:11.832 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:06:11.832 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:06:11.833 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:06:11.833 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic1']\n",
      "23-06-26 08:06:11.834 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:06:11.834 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:06:11.834 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "topic1 = \"test_topic1\"\n",
    "topic2 = \"test_topic2\"\n",
    "sent_msgs_1 = [(f\"msg{i}\" + topic1).encode(\"UTF-8\") for i in range(32)]\n",
    "sent_msgs_2 = [(f\"msg{i}\" + topic2).encode(\"UTF-8\") for i in range(32)]\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    async with create_consumer_and_producer(auto_offset_reset=\"earliest\") as (\n",
    "        consumer,\n",
    "        producer,\n",
    "    ):\n",
    "        [await producer.send(topic1, msg) for msg in sent_msgs_1]\n",
    "        [await producer.send(topic2, msg) for msg in sent_msgs_2]\n",
    "\n",
    "        consumer.subscribe([topic1])\n",
    "        received = await consumer.getmany()\n",
    "        received_msgs = [msg.value for _, msgs in received.items() for msg in msgs]\n",
    "\n",
    "    assert checkEqual(sent_msgs_1, received_msgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bb1c5c5",
   "metadata": {},
   "source": [
    "Check if msgs are received only after subscribing when auto_offset_reset is set to \"latest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6bba51",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:06:17.222 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:06:17.223 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:06:17.223 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:06:17.224 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:06:17.224 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:06:17.225 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:06:17.226 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:06:17.226 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:06:17.227 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "topic = \"test_topic\"\n",
    "sent_msgs_before = [f\"msg{i}\".encode(\"UTF-8\") for i in range(32)]\n",
    "sent_msgs_after = [f\"msg{i}\".encode(\"UTF-8\") for i in range(32, 64)]\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    async with create_consumer_and_producer() as (consumer, producer):\n",
    "        [await producer.send(topic, msg) for msg in sent_msgs_before]\n",
    "\n",
    "        consumer.subscribe([topic])\n",
    "        received = await consumer.getmany()\n",
    "        [await producer.send(topic, msg) for msg in sent_msgs_after]\n",
    "        received = await consumer.getmany()\n",
    "        received_msgs = [msg.value for _, msgs in received.items() for msg in msgs]\n",
    "\n",
    "    assert checkEqual(sent_msgs_after, received_msgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83e9f43a",
   "metadata": {},
   "source": [
    "Check two consumers different groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2db9c2ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:06:42.343 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:06:42.344 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:06:42.345 [ERROR] asyncio: Unclosed AIOKafkaConsumer\n",
      "consumer: <aiokafka.consumer.consumer.AIOKafkaConsumer object>\n",
      "23-06-26 08:06:42.345 [ERROR] asyncio: Unclosed AIOKafkaProducer\n",
      "producer: <aiokafka.producer.producer.AIOKafkaProducer object>\n",
      "23-06-26 08:06:42.345 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:06:42.346 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:06:42.346 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:06:42.347 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:06:42.347 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:06:42.347 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:06:42.348 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:06:42.349 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:06:42.350 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:06:42.350 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:06:42.351 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "topic = \"test_topic\"\n",
    "sent_msgs = [f\"msg{i}\".encode(\"UTF-8\") for i in range(32)]\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    consumer1 = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "        auto_offset_reset=\"earliest\"\n",
    "    )\n",
    "    consumer2 = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "        auto_offset_reset=\"earliest\"\n",
    "    )\n",
    "    producer = fastkafka._aiokafka_imports.AIOKafkaProducer()\n",
    "\n",
    "    await consumer1.start()\n",
    "    await consumer2.start()\n",
    "    await producer.start()\n",
    "\n",
    "    [await producer.send(topic, msg) for msg in sent_msgs]\n",
    "\n",
    "    consumer1.subscribe([topic])\n",
    "    received1 = await consumer1.getmany()\n",
    "\n",
    "    consumer2.subscribe([topic])\n",
    "    received2 = await consumer2.getmany()\n",
    "\n",
    "    received_msgs1 = [msg.value for _, msgs in received1.items() for msg in msgs]\n",
    "    received_msgs2 = [msg.value for _, msgs in received2.items() for msg in msgs]\n",
    "\n",
    "    await consumer1.stop()\n",
    "    await consumer2.stop()\n",
    "    await producer.stop()\n",
    "\n",
    "    assert checkEqual(sent_msgs, received_msgs1), received_msgs1\n",
    "    assert checkEqual(sent_msgs, received_msgs2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1687740a",
   "metadata": {},
   "source": [
    "Check two consumers same group"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c7ed5e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:07:00.711 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:07:00.711 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:07:00.712 [ERROR] asyncio: Unclosed AIOKafkaConsumer\n",
      "consumer: <aiokafka.consumer.consumer.AIOKafkaConsumer object>\n",
      "23-06-26 08:07:00.712 [ERROR] asyncio: Unclosed AIOKafkaProducer\n",
      "producer: <aiokafka.producer.producer.AIOKafkaProducer object>\n",
      "23-06-26 08:07:00.713 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:07:00.713 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:07:00.713 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:07:00.714 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:07:00.714 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:07:00.715 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:07:00.715 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:07:00.715 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:07:00.716 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:07:00.716 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:07:00.716 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "topic = \"test_topic\"\n",
    "sent_msgs = [f\"msg{i}\".encode(\"UTF-8\") for i in range(32)]\n",
    "\n",
    "with InMemoryBroker(num_partitions=5) as broker:\n",
    "    consumer1 = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "        group_id=\"my_group\", auto_offset_reset=\"earliest\"\n",
    "    )\n",
    "    consumer2 = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "        group_id=\"my_group\", auto_offset_reset=\"earliest\"\n",
    "    )\n",
    "    producer = fastkafka._aiokafka_imports.AIOKafkaProducer()\n",
    "\n",
    "    await consumer1.start()\n",
    "    await consumer2.start()\n",
    "    await producer.start()\n",
    "\n",
    "    [await producer.send(topic, msg) for msg in sent_msgs]\n",
    "\n",
    "    consumer1.subscribe([topic])\n",
    "    consumer2.subscribe([topic])\n",
    "\n",
    "    received1 = await consumer1.getmany()\n",
    "    received2 = await consumer2.getmany()\n",
    "\n",
    "    received_msgs1 = [msg.value for _, msgs in received1.items() for msg in msgs]\n",
    "    received_msgs2 = [msg.value for _, msgs in received2.items() for msg in msgs]\n",
    "\n",
    "    await consumer1.stop()\n",
    "    await consumer2.stop()\n",
    "    await producer.stop()\n",
    "\n",
    "    assert checkEqual(sent_msgs, received_msgs1 + received_msgs2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "470875ee",
   "metadata": {},
   "source": [
    "Check for different bootstrap servers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e046307",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23-06-26 08:07:14.565 [INFO] __main__: InMemoryBroker._patch_consumers_and_producers(): Patching consumers and producers!\n",
      "23-06-26 08:07:14.566 [INFO] __main__: InMemoryBroker starting\n",
      "23-06-26 08:07:14.567 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:07:14.567 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:07:14.568 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:07:14.569 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:07:14.570 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:07:14.570 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:07:14.570 [INFO] __main__: AIOKafkaConsumer patched start() called()\n",
      "23-06-26 08:07:14.571 [INFO] __main__: AIOKafkaProducer patched start() called()\n",
      "23-06-26 08:07:14.572 [INFO] __main__: AIOKafkaConsumer patched subscribe() called\n",
      "23-06-26 08:07:14.572 [INFO] __main__: AIOKafkaConsumer.subscribe(), subscribing to: ['test_topic']\n",
      "23-06-26 08:07:14.572 [INFO] __main__: AIOKafkaConsumer patched stop() called\n",
      "23-06-26 08:07:14.573 [INFO] __main__: AIOKafkaProducer patched stop() called\n",
      "23-06-26 08:07:14.573 [INFO] __main__: InMemoryBroker stopping\n"
     ]
    }
   ],
   "source": [
    "topic = \"test_topic\"\n",
    "sent_msgs = [f\"msg{i}\".encode(\"UTF-8\") for i in range(32)]\n",
    "\n",
    "with InMemoryBroker() as broker:\n",
    "    for server in [\"localhost:9092\", \"kafka.airt.ai\"]:\n",
    "        consumer = fastkafka._aiokafka_imports.AIOKafkaConsumer(\n",
    "            bootstrap_servers=server, auto_offset_reset=\"earliest\"\n",
    "        )\n",
    "\n",
    "        producer = fastkafka._aiokafka_imports.AIOKafkaProducer(\n",
    "            bootstrap_servers=server\n",
    "        )\n",
    "\n",
    "        await consumer.start()\n",
    "        await producer.start()\n",
    "\n",
    "        [await producer.send(topic, msg) for msg in sent_msgs]\n",
    "\n",
    "        consumer.subscribe([topic])\n",
    "        received = await consumer.getmany()\n",
    "\n",
    "        received_msgs = [msg.value for _, msgs in received.items() for msg in msgs]\n",
    "\n",
    "        await consumer.stop()\n",
    "        await producer.stop()\n",
    "\n",
    "        assert checkEqual(sent_msgs, received_msgs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1f760f8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
